import numpy as np
import matplotlib.pyplot as plt
import matplotlib

# 设置 Matplotlib 显示中文
matplotlib.rcParams['font.sans-serif'] = ['WenQuanYi Micro Hei']  # SimHei 是黑体，可换成你的系统字体
matplotlib.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 生成训练轮数
epochs = np.arange(1, 1501)

# 模拟损失下降曲线（假设趋势）
total_loss = np.exp(-epochs / 500) * 0.5 + 0.05  # 总损失
classification_loss = np.exp(-epochs / 600) * 0.4 + 0.02  # 分类损失
reconstruction_loss = np.exp(-epochs / 700) * 0.3 + 0.03  # 重建损失

# 绘制训练损失曲线
plt.figure(figsize=(8, 5))
plt.plot(epochs, total_loss, label="总损失 (Total Loss)", color='red')
plt.plot(epochs, classification_loss, label="分类损失 (Classification Loss)", color='blue', linestyle='--')
plt.plot(epochs, reconstruction_loss, label="重建损失 (Reconstruction Loss)", color='green', linestyle='-.')

# 添加标题和标签
plt.xlabel("训练轮数 (Epochs)")
plt.ylabel("损失值 (Loss)")
plt.title("CAE-DDPM 训练损失下降曲线")
plt.legend()
plt.grid(True)

# 显示图表
plt.show()
